package com.davon.risk.controll;

import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.state.*;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.triggers.Trigger;
import org.apache.flink.streaming.api.windowing.triggers.TriggerResult;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import java.time.Duration;
import java.util.Arrays;
import java.util.Iterator;
import java.util.Map;

public class TriggerTest {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

//        SingleOutputStreamOperator<WordAndTs> socketDS = env.socketTextStream("localhost", 9999, "\n")
        SingleOutputStreamOperator<WordAndTs> socketDS = env.fromCollection(Arrays.asList("hive 1 10 0 1001", "flink 2 15 0 1001", "hive 3 10 1 1001", "java 4 100 0 1001", "flink 5 15 1 1001"))
                .map(line -> {
                    String[] fields = line.split(" ");
                    return new WordAndTs(fields[0], Long.valueOf(fields[1]), Long.valueOf(fields[2]), Integer.valueOf(fields[3]), Integer.valueOf(fields[4]));
                })
                .assignTimestampsAndWatermarks(WatermarkStrategy.<WordAndTs>forBoundedOutOfOrderness(Duration.ofSeconds(5L)).withTimestampAssigner(new SerializableTimestampAssigner<WordAndTs>() {
                    @Override
                    public long extractTimestamp(WordAndTs element, long recordTimestamp) {
                        return element.ts * 1000;
                    }
                }));

        SingleOutputStreamOperator<WordCount> wordCntDS = socketDS.keyBy(wordAndTs -> wordAndTs.id)
                .window(TumblingEventTimeWindows.of(Time.days(1L))) // 一天滚动窗口
                .trigger(new MyCateRuleTrigger()) // 自定义 trigger, 每来一条计算一次
                .process(new MyProcessWindowFunc());

        wordCntDS.map(wordCnt -> ">>>>" + wordCnt.word + ":" + wordCnt.isRefund)
                .print();

        env.execute();
    }


    // 自定义 Trigger
    static class MyCateRuleTrigger extends Trigger<WordAndTs, TimeWindow> {

        @Override
        public TriggerResult onElement(WordAndTs element, long timestamp, TimeWindow window, TriggerContext ctx) throws Exception {
            // 注册 EventTimeTimer，窗口结束时触发
            ctx.registerEventTimeTimer(window.getEnd());
            // 每来一条都 计算并清除，所以在后面的 ProcessWindowFunc 中，每次迭代器里面只有一个元素
            // PURGE 清除 只清除窗口的状态 如：窗口已经拿到的元素集合，
            // 不会清除后面的窗口处理函数(本例中的MyProcessWindowFunc) 中的状态
            return TriggerResult.FIRE_AND_PURGE;
        }

        @Override
        public TriggerResult onProcessingTime(long time, TimeWindow window, TriggerContext ctx) throws Exception {
            return TriggerResult.CONTINUE;
        }

        @Override
        public TriggerResult onEventTime(long time, TimeWindow window, TriggerContext ctx) throws Exception {
            // EventTimeTimer 计算并清除
            return TriggerResult.FIRE_AND_PURGE;
        }

        @Override
        public void clear(TimeWindow window, TriggerContext ctx) throws Exception {
        }
    }

    //自定义 ProcessWindowFunction
    static class MyProcessWindowFunc extends ProcessWindowFunction<WordAndTs, WordCount, Integer, TimeWindow> {
        // 针对阈值被调小的情况，可能需要再增加一个valueState指示mapState中的数据是否已经输出，
        // 可能需要进行撤回,此处未考虑此情况

        ValueStateDescriptor<Long> valueState;	// 存储累计 rewared
        MapStateDescriptor<String, Boolean> mapState;	// 存储所有的已支付订单
        MapStateDescriptor<String, Boolean> unmatchedRefundMap;	// 存储乱序导致的先到的退单订单
        private final Integer rewardLimit = 20;	// 简化的固定rewared 阈值，生产中从 input element 中得到

        @Override
        public void open(Configuration parameters) throws Exception {
            super.open(parameters);
            // 初始化状态，设置 ttl
            valueState = new ValueStateDescriptor<>("word_cnt", Long.class);
            mapState = new MapStateDescriptor<>("word_set", String.class, Boolean.class);
            unmatchedRefundMap = new MapStateDescriptor<>("unmatched_refund_set", String.class, Boolean.class);

            StateTtlConfig ttlConfig = StateTtlConfig
                    .newBuilder(org.apache.flink.api.common.time.Time.seconds(1))
                    .setUpdateType(StateTtlConfig.UpdateType.OnCreateAndWrite)
                    .setStateVisibility(StateTtlConfig.StateVisibility.NeverReturnExpired)
                    .build();
            valueState.enableTimeToLive(ttlConfig);
            mapState.enableTimeToLive(ttlConfig);
            unmatchedRefundMap.enableTimeToLive(ttlConfig);
        }

        @Override
        public void process(Integer s, Context context, Iterable<WordAndTs> elements, Collector<WordCount> out) throws Exception {
            WordAndTs wordAndTs = elements.iterator().next();
            ValueState<Long> rewardAmt = context.windowState().getState(valueState);
            Long rewardAmtVal = rewardAmt.value();
            if (rewardAmtVal == null) {
                // valueState.value 初始是 null
                rewardAmtVal = 0L;
            }
//            System.out.println("before rewardAmtVal: " + rewardAmtVal);
            System.out.println("current word: " + wordAndTs.word + "\trefund: " + wordAndTs.isRefund);

            MapState<String, Boolean> beforeOrders = context.windowState().getMapState(mapState);
            MapState<String, Boolean> unmatchedRefundOrders = context.windowState().getMapState(unmatchedRefundMap);
            Iterator<Map.Entry<String, Boolean>> ite = beforeOrders.iterator();

            if (wordAndTs.isRefund == 0) {	// 非退单订单
                // 检查该订单是否已有先到的退单
                if (unmatchedRefundOrders.contains(wordAndTs.word)) {
                    unmatchedRefundOrders.remove(wordAndTs.word);
                    return;
                }
                // 对满足的订单输出
                if (rewardAmtVal >= rewardLimit) {
                    out.collect(new WordCount(wordAndTs.word, 0));
                } else if (rewardAmtVal + wordAndTs.reward >= rewardLimit){
                    out.collect(new WordCount(wordAndTs.word, 0));
                    while (ite.hasNext()) {
                        out.collect(new WordCount(ite.next().getKey(), 0));
                    }
                }
                // 更新状态
                rewardAmtVal +=  wordAndTs.reward;
                rewardAmt.update(rewardAmtVal);
                beforeOrders.put(wordAndTs.word, null);


            } else if (wordAndTs.isRefund == 1) {	// 退单订单
                // 是否是先到的
                if (!beforeOrders.contains(wordAndTs.word)) {
                    unmatchedRefundOrders.put(wordAndTs.word, null);
                    return;
                }

                // 对退单造成的不满足阈值进行召回
                if (rewardAmtVal >= rewardLimit && rewardAmtVal - wordAndTs.reward < rewardLimit) {
                    while (ite.hasNext()) {
                        out.collect(new WordCount(ite.next().getKey(), 1));
                    }
                } else if (rewardAmtVal - wordAndTs.reward >= rewardLimit){
                    out.collect(new WordCount(wordAndTs.word, 1));
                }
                // 更新状态
                rewardAmtVal -=  wordAndTs.reward;
                rewardAmt.update(rewardAmtVal);
                beforeOrders.remove(wordAndTs.word);
            }
//            System.out.println("after rewardAmtVal: " + rewardAmtVal);
        }
    }
}

class WordAndTs{
    public String word;
    public Long ts;
    public Long reward;
    public Integer isRefund;
    public Integer id;

    public WordAndTs(String word, Long ts, Long reward, Integer isRefund, Integer id) {
        this.id = id;
        this.word = word;
        this.ts = ts;
        this.reward = reward;
        this.isRefund = isRefund;
    }
}

class WordCount {
    public String word;
    public Integer isRefund;

    public WordCount(String word, Integer isRefund) {
        this.word = word;
        this.isRefund = isRefund;
    }
}
